// Copyright 2022 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef COMPONENTS_OMNIBOX_BROWSER_AUTOCOMPLETE_SCORING_MODEL_SERVICE_H_
#define COMPONENTS_OMNIBOX_BROWSER_AUTOCOMPLETE_SCORING_MODEL_SERVICE_H_

#include <memory>

#include "base/functional/callback.h"
#include "base/memory/scoped_refptr.h"
#include "base/task/cancelable_task_tracker.h"
#include "base/task/sequenced_task_runner.h"
#include "components/keyed_service/core/keyed_service.h"
#include "components/omnibox/browser/autocomplete_scoring_model_executor.h"
#include "components/omnibox/browser/autocomplete_scoring_model_handler.h"
#include "components/optimization_guide/core/optimization_guide_model_provider.h"
#include "third_party/abseil-cpp/absl/types/optional.h"
#include "third_party/metrics_proto/omnibox_event.pb.h"

// Autocomplete scoring service using machine learning models via
// OptimizationGuide's model handler.
class AutocompleteScoringModelService : public KeyedService {
 public:
  using ResultCallback =
      base::OnceCallback<void(std::tuple<absl::optional<float>, size_t, GURL>)>;
  using ScoringSignals =
      ::metrics::OmniboxEventProto::Suggestion::ScoringSignals;

  explicit AutocompleteScoringModelService(
      optimization_guide::OptimizationGuideModelProvider* model_provider);
  ~AutocompleteScoringModelService() override;

  // Disallow copy/assign.
  AutocompleteScoringModelService(const AutocompleteScoringModelService&) =
      delete;
  AutocompleteScoringModelService& operator=(
      const AutocompleteScoringModelService&) = delete;

  // Invokes the model to score the given `scoring_signals` and calls
  // `result_callback` with an optional relevance score generated by the model.
  void ScoreAutocompleteUrlMatch(base::CancelableTaskTracker* tracker,
                                 const ScoringSignals& scoring_signals,
                                 size_t match_index,
                                 GURL match_destination_url,
                                 ResultCallback result_callback);

  // Returns whether the scoring model is loaded and the pointer to the
  // handler is non-null.
  bool UrlScoringModelAvailable();

 private:
  // Processes the model output and invokes the callback with the relevance
  // score from the model output. Invokes the callback with nullopt if the model
  // output is nullopt or an empty vector (which is unexpected).
  void ProcessModelOutput(
      ResultCallback result_callback,
      size_t match_index,
      GURL match_destination_url,
      const absl::optional<AutocompleteScoringModelExecutor::ModelOutput>&
          model_output);

  scoped_refptr<base::SequencedTaskRunner> model_executor_task_runner_;

  // Autocomplete URL scoring model.
  std::unique_ptr<AutocompleteScoringModelHandler> url_scoring_model_handler_ =
      nullptr;
};

#endif  // COMPONENTS_OMNIBOX_BROWSER_AUTOCOMPLETE_SCORING_MODEL_SERVICE_H_
